package io.gatling.extensions.kafka.protocol

import java.util
import java.util.Properties

import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.gatling.core.{CoreComponents, protocol}
import org.apache.kafka.clients.producer.KafkaProducer
import org.apache.kafka.common.serialization.StringSerializer

case class KafkaProducerProtocol(servers: String, acks: String = "1") extends Protocol {
  type Components = KafkaProducerComponents

  val props: Properties = new Properties
  props.put("bootstrap.servers", servers)
  props.put("acks", acks)
  props.put("retries", "3")
  props.put("batch.size", "10000")
  props.put("linger.ms", "1000")
  props.put("request.timeout.ms", "1000")
  props.put("buffer.memory", "33554432")
  props.put("key.serializer", classOf[StringSerializer])
  props.put("value.serializer", classOf[StringSerializer])
}

object KafkaProducerProtocol {

  val kafkaProtocolKey: ProtocolKey[KafkaProducerProtocol, KafkaProducerComponents] = new ProtocolKey[KafkaProducerProtocol, KafkaProducerComponents] {

    override def protocolClass: Class[protocol.Protocol] = classOf[KafkaProducerProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): KafkaProducerProtocol = KafkaProducerProtocol("192.168.106.130:9092,192.168.106.131:9092,192.168.106.132:9092")

    override def newComponents(coreComponents: CoreComponents): KafkaProducerProtocol => KafkaProducerComponents = {
      kafkaProducerProtocol => KafkaProducerComponents(kafkaProducerProtocol)
    }
  }
}

case class KafkaProducerComponents(kafkaProducerProtocol: KafkaProducerProtocol) extends ProtocolComponents {

  val producerMap: util.HashMap[Long, KafkaProducer[String, String]] = new util.HashMap[Long, KafkaProducer[String, String]]()

  override def onStart: Session => Session = session => {
    producerMap.put(session.userId, new KafkaProducer[String, String](kafkaProducerProtocol.props))
    session
  }

  override def onExit: Session => Unit = session => {
    producerMap.get(session.userId).close()
  }
}

case class KafkaProducerProtocolBuilder(servers: String, acks: String = "1") {
  def build: KafkaProducerProtocol = KafkaProducerProtocol(servers, acks)
}

object KafkaProducerProtocolBuilder {
  implicit def toKafkaProducerProtocol(builder: KafkaProducerProtocolBuilder): KafkaProducerProtocol = builder.build
}
